package storage

import (
	"context"

	"gitcode.com/eteam/sae.git/types"
	"go.mongodb.org/mongo-driver/v2/bson"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

type ArtifactInfoDbItem struct {
	ArtifactId  string       `bson:"_id"`
	TaskId      string       `bson:"taskId"`
	Name        string       `bson:"name"`
	Description string       `bson:"description"`
	TimeStamp   int64        `bson:"timeStamp"`
	Parts       []types.Part `bson:"parts"`
}

func (item *ArtifactInfoDbItem) ToArtifact() *types.Artifact {
	if item == nil {
		return nil
	}
	return &types.Artifact{
		ArtifactId:  item.ArtifactId,
		Name:        item.Name,
		Description: item.Description,
		TimeStamp:   item.TimeStamp,
		Parts:       item.Parts,
	}
}

type _ArtifactInfoDao struct {
	mc *mongo.Collection
}

func newArtifactInfoDao(db *mongo.Database) (*_ArtifactInfoDao, error) {
	mc := db.Collection("artifact_info")
	_, err := mc.Indexes().CreateOne(context.Background(), mongo.IndexModel{
		Keys: bson.D{{"taskId", 1}, {"timeStamp", 1}},
	})
	if err != nil {
		return nil, err
	}
	return &_ArtifactInfoDao{mc: mc}, nil
}

func (dao *_ArtifactInfoDao) Insert(ctx context.Context, dbItem *ArtifactInfoDbItem) error {
	_, err := dao.mc.InsertOne(ctx, dbItem)
	return err
}

func (dao *_ArtifactInfoDao) ListByTask(ctx context.Context, taskIdList []string) ([]*ArtifactInfoDbItem, error) {
	if len(taskIdList) == 0 {
		return []*ArtifactInfoDbItem{}, nil
	}
	cursor, err := dao.mc.Find(ctx, bson.M{"taskId": bson.M{"$in": taskIdList}},
		options.Find().SetSort(bson.M{"timeStamp": 1}))
	if err != nil {
		return nil, err
	}
	retList := []*ArtifactInfoDbItem{}
	err = cursor.All(ctx, &retList)
	return retList, err
}

func (dao *_ArtifactInfoDao) Update(ctx context.Context, dbItem *ArtifactInfoDbItem) error {
	_, err := dao.mc.UpdateOne(ctx, bson.M{"_id": dbItem.ArtifactId}, bson.M{"$set": dbItem})
	return err
}
